package com.bruce.ai.alibaba.graph.controller;

import java.util.HashMap;
import java.util.Map;

import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.OverAllState;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;

@RestController
@RequestMapping("/graph/stream")
public class StreamController {

    private static final Logger logger = LoggerFactory.getLogger(StreamController.class);

    private final CompiledGraph compiledGraph;

    public StreamController(@Qualifier("streamGraph")StateGraph stateGraph) throws GraphStateException {
        this.compiledGraph = stateGraph.compile();
    }

    @GetMapping(value = "/expand", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？", required = false) String query,
                                                @RequestParam(value = "expander_number", defaultValue = "3", required = false) Integer  expanderNumber,
                                                @RequestParam(value = "thread_id", defaultValue = "bruce_sg", required = false) String threadId) throws GraphRunnerException {
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("query", query);
        objectMap.put("expander_number", expanderNumber);

        return Mono.defer(() -> Mono.justOrEmpty(compiledGraph.invoke(objectMap, runnableConfig))
                .map(OverAllState::data)
                .map(data -> ServerSentEvent.builder(data.toString()).build()))
            .flux()
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }

    /**
     * 测试异常处理的端点 - 仅用于演示 doOnError 触发
     */
    @GetMapping(value = "/test-error", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> testError() {
        return Flux.just("data1", "data2", "error")
                .delayElements(java.time.Duration.ofSeconds(1))
                .map(data -> {
                    if ("error".equals(data)) {
                        throw new RuntimeException("模拟的流处理异常");
                    }
                    return ServerSentEvent.builder(data).build();
                })
                .doOnCancel(() -> logger.info("Test: Client disconnected from stream"))
                .doOnError(e -> logger.error("Test: Error occurred during streaming", e));
    }

}
